#!/home/robot/anaconda3/envs/yolograsp/bin/python3
import os
import numpy as np
import sys
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '../'))
from mmdet.apis import inference_detector, init_detector, show_result_pyplot
sys.path.pop(0)
import warnings
warnings.filterwarnings('ignore')
import rospy
from vision_messages.srv import InstanceMask, InstanceMaskResponse

# yjh
sys.path.insert(1, os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../..')) # src/
from utilities.tools.image_converter import convert_msg_to_nparray
sys.path.pop(1)


cat_id = 41
# cat       coco_id     nocs_id
# cup       41          5
# bowl      45          1
# bottle    39          0
coco_list = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
            'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
            'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
            'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
            'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
            'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
            'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
            'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
            'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
            'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
            'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
            'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
            'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
            'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
id_map = {
    41: 5,
    45: 1,
    39: 0
}

class Solo(object):
    def __init__(self):
        # yjh
        config=os.path.join(os.path.dirname(os.path.abspath(__file__)), '../configs/solo/decoupled_solo_light_r50_fpn_3x_coco.py')
        checkpoint=os.path.join(os.path.dirname(os.path.abspath(__file__)), '../decoupled_solo_weights.pth')
        device='cuda:0'
        self.score_thr=0.6
        self.model=init_detector(config, checkpoint, device)

        rospy.init_node('mask_seg', anonymous=True)
        rospy.Service("/perception/mmdseg/solo", InstanceMask, self.segmentation_callback)
        rospy.wait_for_service("/perception/mmdseg/solo")
        rospy.loginfo("Solo init done! Ready to get image!")

    def segmentation_callback(self, req):
        rospy.loginfo("Get image!")
        mask = None
        # img, model, args = req.img, req.model, req.args
        # mask = self.get_mask(img, model, args)

        # yjh
        img = convert_msg_to_nparray(req.color)
        result = self.get_mask(img)
        # Here, we only see if we can get category:
        cate = result['class_ids']
        rospy.loginfo("Return Mask Results!")
        return InstanceMaskResponse(cate)
    
    def get_bbox_from_mask(self, mask):
        # mask: bool array, shape: (H, W)
        # return: (x1, y1, x2, y2)
        mask_ids = np.argwhere(mask == True)
        left_top = mask_ids.min(axis=0)
        right_bottom = mask_ids.max(axis=0)
        return np.array([left_top[1], left_top[0], right_bottom[1], right_bottom[0]])

    def get_mask(self, img, show_mask=False):
        # if model is None:
        #     model = init_detector(args.config, args.checkpoint, device=args.device)
        result_ins = {}
        result = inference_detector(self.model, img)
        # show the results
        if show_mask:
            show_result_pyplot(self.model, img, result, score_thr=self.score_thr)
        # Here, we only consider the scene with one object.
        max_score = 0.
        for id in id_map.keys():
            if result[0][id].shape[0] > 0:
                score = result[0][id][:, 4].max()
                if score > max_score:
                    max_score = score
                    cat_id = id
        result_ins['class_name'] = coco_list[cat_id]
        result_ins['class_ids'] = id_map[cat_id]
        result_ins['scores'] = result[0][cat_id][:, 4]
        best_idx = result[0][cat_id][:, 4].argmax()
        result_ins['scores'] = result_ins['scores'][best_idx]
        result_ins['masks'] = np.array(result[1][cat_id])[best_idx]
        # result_ins['rois'] = result[0][cat_id][best_idx, :4]
        result_ins['rois'] = self.get_bbox_from_mask(result_ins['masks'])
        return result_ins


if __name__ == '__main__':
    seg_srv = Solo()
    rospy.spin()
